from pox.core import core
from pox.lib.util import dpid_to_str
import pox.openflow.libopenflow_01 as of
from pox.lib.recoco import Timer
from pytrie import SortedStringTrie as trie

log = core.getLogger()

class MyComponent(object):
        def __init__(self):
            core.openflow.addListeners(self)
            self.flag = 0
            self.connection = 0
            self.hhhlist = ['x'] #List of HHH's of the previous iteration
            self.new_hhhlist = [] # List of newly identified HHH's
            self.trie = trie(x0=0,x1=0)
            self.link_cap = 10.0 #If node weight is greater than this value , its a heavy hitter

        def handle_timer_elapse(self):
            reqmsg = of.ofp_stats_request(type = of.OFPST_FLOW, body = of.ofp_flow_stats_request())
            self.connection.send(reqmsg)

        def install_flow(self,install,priority = 24 ,src_ip="10.0.0.0/24",dst_ip="10.0.0.3"):
            msg = of.ofp_flow_mod()
            msg.match.nw_src = src_ip
            msg.match.nw_dst = dst_ip
            msg.priority = priority
            if install == False:
                msg.command = of.OFPFC_DELETE
            else:
                msg.command = of.OFPFC_ADD
            msg.match.dl_type = 0x800 # Type for IPv4
#            msg.match.tp_dst = 80
            msg.actions.append(of.ofp_action_output(port = 3))
            self.connection.send(msg)
            msg.match.dl_type = 0x0806 # Type for ARP messages , so it seems
            self.connection.send(msg)
            msg.match.dl_type = 0x800
            msg.match.nw_src = dst_ip # Install reverse flow and ARP entry as well.
            msg.match.nw_dst = src_ip
            msg.actions.append(of.ofp_action_output(port = of.OFPP_FLOOD))
            self.connection.send(msg)
            msg.match.dl_type = 0x0806
            self.connection.send(msg)

        def formatIpAddr(self, ipaddr):
            splitaddr = ipaddr.split('.')
            intaddr = int(splitaddr[3])
            binary_addr = "{0:b}".format(intaddr)
            for i in range(0,8-len(binary_addr)):
                binary_addr = "0" + binary_addr
            binary_addr = "x" + binary_addr[:4] #We only need the higher nibble here
            return binary_addr

        def IPifier(self,prefix,ip="10.0.0."): # takes in string of the form "x00" and gives out an IP of the form 10.0.0.something/mask , assuming only last 8 bits of IP are under variation
            clean_prefix = prefix[1:]
            mask = 24+ len(clean_prefix)
            for i in range(0,4-len(clean_prefix)):
                clean_prefix = clean_prefix + "0"
            clean_prefix = clean_prefix + "0000" #Since this prefix is the higher nibble in the 8 bit integer
            formed_ip = ip + str(int(clean_prefix,2)) + "/" + str(mask)
            log.debug("IP to be monitored: %s",formed_ip)
            return formed_ip

        def priorityCalculator(self,element): # We use subnet mask as priority here, higher the value , the higher the priority
            element_priority = element[1:]
            priority = 24 + len(element_priority)
            return priority

        def _handle_FlowStatsReceived(self,event):
#            print "FlowStatsReceived"
            print "=================================================================================="
            total_count = 0.0
            for j in event.stats: # Get total packet count
                if j.match.dl_type==0x800:
                    total_count+=j.packet_count

            for f in event.stats:
                log.info("number of packets in previous interval %d , priority: %s, type: %s , ip %s , total count %d",f.packet_count,f.priority,f.match.dl_type, f.match.nw_src, total_count)
                if f.match.dl_type==0x800:
                    formatted_addr = self.formatIpAddr(str(f.match.nw_src))
                    try:
                        found_match = self.trie.longest_prefix(formatted_addr) # TODO: If match not found? Handle exception and insert to tree
                    except Exception:
                        log.debug("Something bad happened")
                        continue
                    if total_count==0.0:  # There was no traffic!
                        flow_percentage =0.0
                    else:
                        flow_percentage = (100 * f.packet_count) / total_count
                    self.trie.__setitem__(found_match,flow_percentage+self.trie.longest_prefix_value(found_match, default='N/A'))
            log.debug("Tree values : %s",self.trie)
            self.getHeavyHitters()
            self.hhhlist = self.new_hhhlist
            self.new_hhhlist = []
            self.resetFlowTables()
            self.trie = trie(x0=0,x1=0) #Reset tree for next iteration
            log.debug("Repopulating with new rules")
            for element in self.hhhlist:
                tup = self.reproduce(element)
                if tup[0] == tup[1] == "0": # If HHH is child , monitor the child in the next iteration
                    self.install_flow(True,self.priorityCalculator(element),self.IPifier(element))
                    self.trie.__setitem__(element,0)
                else: # Monitor the children
                    self.install_flow(True,self.priorityCalculator(tup[0]),self.IPifier(tup[0]))
                    self.install_flow(True,self.priorityCalculator(tup[1]),self.IPifier(tup[1]))
                    self.trie.__setitem__(tup[0],0)
                    self.trie.__setitem__(tup[1],0) # Insert children into the monitoring tree

        def resetFlowTables(self):
            msg = of.ofp_flow_mod(command=of.OFPFC_DELETE)
            for connection in core.openflow.connections:
                connection.send(msg)
            log.debug("Clearing all flows from Switch")
            self.install_flow(True,25,"10.0.0.0/25")
            self.install_flow(True,25,"10.0.0.128/25")

        def reproduce(self,element):
            if len(element)>=5:
                return ("0","0")
            else:
                Lchild = element + "0"
                Rchild = element + "1"
                return (Lchild,Rchild)

        def getHeavyHitters(self):
            for current_pref in self.hhhlist:
                if current_pref.__len__()==5: #This is a leaf
                    parentval = self.trie.longest_prefix_value(current_pref)
                    log.debug("Current pref in getheavyhitters : %s parentval: %s", current_pref,parentval)
                    if parentval > self.link_cap:
                        self.new_hhhlist.append(current_pref)
                        log.debug("In a leaf HHH")
                    continue
                Lchild_pref = current_pref + "0"
                Rchild_pref = current_pref + "1"
                Lchildval = self.trie.longest_prefix_value(Lchild_pref)
                Rchildval = self.trie.longest_prefix_value(Rchild_pref)
                current_val = Lchildval + Rchildval #Current node is the parent, current_val is val of parent
                log.debug("Current pref in getheavyhitters : %s parentval: %s", current_pref,current_val)
                if (current_val > self.link_cap) and (Lchildval < self.link_cap) and (Rchildval < self.link_cap):
                    self.new_hhhlist.append(current_pref)
                elif Lchildval > self.link_cap and Rchildval < self.link_cap:
                    self.new_hhhlist.append(Lchild_pref)
                elif Rchildval > self.link_cap and Lchildval < self.link_cap:
                    self.new_hhhlist.append(Rchild_pref)
                elif Rchildval > self.link_cap and Lchildval > self.link_cap:
                    self.new_hhhlist.append(Lchild_pref)
                    self.new_hhhlist.append(Rchild_pref)
            if len(self.new_hhhlist)==0:
                self.new_hhhlist.append('x')
            log.debug("Heavy hitter list %s", self.new_hhhlist)

        def _handle_ConnectionUp(self,event):
            log.debug("Switch %s has come up.",dpid_to_str(event.dpid))

        def _handle_ConnectionDown(self,event):
            log.debug("Switch went down")

        def _handle_PortStatus(self,event):
            if event.added:
                action = "added"
            elif event.deleted:
                action = "removed"
            else:
                action = "modified"
            print "Port %s on Switch %s has been %s" %(event.port, event.dpid, action)

        def _handle_PacketIn(self,event):
            log.debug("Received packet  port %s ", event.port)
            self.connection = event.connection
            if self.flag==0:
                self.flag = 1
                self.install_flow(True,25,"10.0.0.0/25")
                self.install_flow(True,25,"10.0.0.128/25")
                Timer(10,self.handle_timer_elapse,recurring = True)


def launch():
	core.registerNew(MyComponent)

